import io

from PIL import Image
from flask import Flask, request, jsonify
from tensorflow.keras.applications import VGG16
from tensorflow.keras.applications.vgg16 import preprocess_input
from tensorflow.keras.preprocessing import image
from tensorflow.keras.models import Model
import numpy as np
import os
# 环境设置，用于避免某些 TensorFlow 的警告和错误
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # 设置 TensorFlow 的日志级别
os.environ['KMP_DUPLICATE_LIB_OK']='True' # 环境变量用于解决 OpenMP 在 macOS 上的问题

# 修改工作目录到脚本所在的目录，并在这个目录下创建.keras目录，将其设置为Keras下载模型的目录
script_dir = os.getcwd()
print("依赖目录：",script_dir)
os.makedirs(os.path.join(script_dir, '.keras'), exist_ok=True)
weights_path = os.path.join(script_dir, '.keras', 'vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5')

# 加载预训练的 VGG16 模型，不包括最后的全连接层
base_model = VGG16(weights=weights_path, include_top=False)

# 定义新的模型，将原 VGG16 模型的输出直接作为新模型的输出
model = Model(inputs=base_model.input, outputs=base_model.output)

# 定义 Flask 应用
app = Flask(__name__)

# 定义函数 extract_features，用于从给定图片中提取特征
# 定义函数 extract_features，能够区分输入是文件路径还是文件内容
def extract_features(img_content_or_path, model, is_path=True):
    if is_path:
        img = Image.open(img_content_or_path)
    else:
        img = Image.open(io.BytesIO(img_content_or_path))
    img = img.convert('RGB')
    img = img.resize((224, 224))
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)
    features = model.predict(img_array)
    flatten_features = features.flatten()
    normalized_features = flatten_features / np.linalg.norm(flatten_features)
    return normalized_features

# 定义函数 calculate_similarity，用于计算两个特征之间的相似度
def calculate_similarity(features1, features2):
    similarity = np.dot(features1, features2)
    return similarity

# 定义 Flask 路由，对于通过HTTP请求发送的文件进行处理
@app.route('/calculate_similarity', methods=['POST'])
def calculate_similarity_endpoint():
    if 'image1' not in request.files or 'image2' not in request.files:
        return 'No file part', 400
    file1 = request.files['image1'].read()
    file2 = request.files['image2'].read()
    features1 = extract_features(file1, model, is_path=False)
    features2 = extract_features(file2, model, is_path=False)
    similarity = calculate_similarity(features1, features2)
    return jsonify({"Similarity": similarity.item()})

# 定义 Flask 路由，对于通过文件路径进行处理
@app.route('/calculate_similarity_path', methods=['POST'])
def calculate_similarity_path_endpoint():
    data = request.get_json()
    features1 = extract_features(data['image_path1'], model)
    features2 = extract_features(data['image_path2'], model)
    similarity = calculate_similarity(features1, features2)
    return jsonify({"Similarity": similarity.item()})

# 启动 Flask 服务
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)